{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Graph_regularization_for_Twitter_rumour_veracity_classification_using_natural_graphs.ipynb",
      "provenance": [],
      "private_outputs": true,
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ktxosp9yG6ci",
        "colab_type": "text"
      },
      "source": [
        "##### Copyright 2021 Google LLC\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oSo7RZo4WjRs",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wCliuvuHPi0W",
        "colab_type": "text"
      },
      "source": [
        "# Graph regularization for Twitter rumour veracity classification using natural graphs\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Sy4yBd23A_AN",
        "colab_type": "text"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial uses the PHEME dataset for veracity classification of Twitter rumours, a binary classification task.\n",
        "\n",
        "Tweet texts are used as input for computing embeddings, e.g. with ALBERT. Those representations are then used as features for the baseline MLP classification model, as well as for the graph regularized version, which uses the structure defined by the tweet replies as a natural graph."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a6jJnC6JvaOt",
        "colab_type": "text"
      },
      "source": [
        "## Install and import dependencies"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9pYjAVnf5j6s",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install --quiet neural-structured-learning\n",
        "!pip install --quiet transformers\n",
        "!pip install --quiet tokenizers\n",
        "!pip install --quiet sentencepiece"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mZN6xDms5tEu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "import json\n",
        "import random\n",
        "import pprint\n",
        "\n",
        "import numpy as np\n",
        "\n",
        "import neural_structured_learning as nsl\n",
        "\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "import sentencepiece\n",
        "\n",
        "from tokenizers import BertWordPieceTokenizer\n",
        "\n",
        "from transformers import BertConfig\n",
        "from transformers import BertTokenizer, TFBertModel\n",
        "from transformers import XLNetTokenizer, TFXLNetModel\n",
        "from transformers import RobertaTokenizer, TFRobertaModel\n",
        "from transformers import AlbertTokenizer, TFAlbertModel\n",
        "from transformers import T5Tokenizer, TFT5Model\n",
        "from transformers import ElectraTokenizer, TFElectraModel\n",
        "\n",
        "# Resets notebook state\n",
        "tf.keras.backend.clear_session()\n",
        "\n",
        "print(\"Version: \", tf.__version__)\n",
        "print(\"Eager mode: \", tf.executing_eagerly())\n",
        "print(\n",
        "    \"GPU is\",\n",
        "    \"available\" if tf.config.list_physical_devices(\"GPU\") else \"NOT AVAILABLE\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cbh3CaQZHOZh",
        "colab_type": "text"
      },
      "source": [
        "## Dataset description\n",
        "\n",
        "The PHEME dataset includes the following, which are used in this tutorial. \n",
        "\n",
        "* Twitter rumours associated with diverse news events\n",
        "* Source tweets, each of which is accompanied by reaction tweets\n",
        "* The reply structure for each source and corresponding reaction tweets\n",
        "* Veracity label annotations (by professional journalists) for each source tweet.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9YP96mYyvqz4",
        "colab_type": "text"
      },
      "source": [
        "## Download and extract the dataset\n",
        "\n",
        "### Attribution\n",
        "\n",
        "\n",
        "We use the [PHEME dataset for Rumour Detection and Veracity Classification](https://figshare.com/articles/dataset/PHEME_dataset_for_Rumour_Detection_and_Veracity_Classification/6392078)\n",
        "created by Elena Kochkina, Maria Liakata, and Arkaitz Zubiaga.\n",
        "\n",
        "This data set is licensed under [CC BY 4.0](https://creativecommons.org/licenses/by/4.0/).\n",
        "The underlying data is used without modifications as labels or for computing model features."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GTRXhIgB57d5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!mkdir /tmp/PHEME\n",
        "!wget --quiet -P /tmp/PHEME https://ndownloader.figshare.com/files/11767817"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PAQiiuKb4A68",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!ls -l /tmp/PHEME"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Kgl2GBLmv--4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!tar -C /tmp/PHEME -xvzf /tmp/PHEME/11767817"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PTyvV37PuaBw",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!ls /tmp/PHEME/all-rnr-annotated-threads/"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H_WgmY-RvxOU",
        "colab_type": "text"
      },
      "source": [
        "## Convert rumor annotations to labels\n",
        "\n",
        "Code available in `convert_veracity_annotations.py` together with the [PHEME dataset](https://figshare.com/articles/PHEME_dataset_for_Rumour_Detection_and_Veracity_Classification/6392078)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TFX9lTDBAwJF",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\"\"\"\n",
        "Python 3 function to convert rumour annotations into True, False, Unverified\n",
        "\"\"\"\n",
        "\n",
        "def convert_annotations(annotation, string = True):\n",
        "    if 'misinformation' in annotation.keys() and 'true'in annotation.keys():\n",
        "        if int(annotation['misinformation'])==0 and int(annotation['true'])==0:\n",
        "            if string:\n",
        "                label = \"unverified\"\n",
        "            else:\n",
        "                label = 2\n",
        "        elif int(annotation['misinformation'])==0 and int(annotation['true'])==1 :\n",
        "            if string:\n",
        "                label = \"true\"\n",
        "            else:\n",
        "                label = 1\n",
        "        elif int(annotation['misinformation'])==1 and int(annotation['true'])==0 :\n",
        "            if string:\n",
        "                label = \"false\"\n",
        "            else:\n",
        "                label = 0\n",
        "        elif int(annotation['misinformation'])==1 and int(annotation['true'])==1:\n",
        "            print (\"OMG! They both are 1!\")\n",
        "            print(annotation['misinformation'])\n",
        "            print(annotation['true'])\n",
        "            label = None\n",
        "            \n",
        "    elif 'misinformation' in annotation.keys() and 'true' not in annotation.keys():\n",
        "        # all instances have misinfo label but don't have true label\n",
        "        if int(annotation['misinformation'])==0:\n",
        "            if string:\n",
        "                label = \"unverified\"\n",
        "            else:\n",
        "                label = 2\n",
        "        elif int(annotation['misinformation'])==1:\n",
        "            if string:\n",
        "                label = \"false\"\n",
        "            else:\n",
        "                label = 0\n",
        "                \n",
        "    elif 'true' in annotation.keys() and 'misinformation' not in annotation.keys():\n",
        "        print ('Has true not misinformation')\n",
        "        label = None\n",
        "    else:\n",
        "        print('No annotations')\n",
        "        label = None\n",
        "           \n",
        "    return label"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ht-VJz75v9tG",
        "colab_type": "text"
      },
      "source": [
        "## Load annotated rumours from dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nZSeSj0b73gc",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def load_pheme_data(parent_dir):\n",
        "  \"\"\" Loads and returns tweets and their replies from input `parent_dir`.\n",
        "  \n",
        "  Args:\n",
        "    parent_dir: A string with full path to directory where the data lies in.\n",
        "\n",
        "  Returns:\n",
        "    A tuple T (annotation, structure, source_tweets, reactions) such that:\n",
        "      Each of them is a dictionary directly read from the underlying JSON\n",
        "      structure provided in the PHEME dataset.\n",
        "  \"\"\"\n",
        "  with open(parent_dir + '/annotation.json') as f:\n",
        "    annotation = json.load(f)\n",
        "  with open(parent_dir + '/structure.json') as f:\n",
        "    structure = json.load(f)\n",
        "  source_tweets = {}\n",
        "  for f in os.scandir(parent_dir + '/source-tweets'):\n",
        "    if f.name[0] != '.':\n",
        "      with open(f.path) as json_file:\n",
        "        source_tweets[f.name.split('.json')[0]] = json.load(json_file)\n",
        "  reactions = {}\n",
        "  for f in os.scandir(parent_dir + '/reactions'):\n",
        "    if f.name[0] != '.':\n",
        "      with open(f.path) as json_file:\n",
        "        reactions[f.name.split('.json')[0]] = json.load(json_file)\n",
        "  return annotation, structure, source_tweets, reactions"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rjmKC0to53MX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def load_labels_and_texts(topics):\n",
        "  \"\"\"Reads verified rumour tweets, replies and labels for input `topics`.\n",
        "  Non-rumours or unverified tweet threads aren't included in returned dataset.\n",
        "\n",
        "  Args:\n",
        "    topics: A List of strings, each containing the full path to a topic\n",
        "    to be read.\n",
        "\n",
        "  Returns:\n",
        "    A List of dictionaries such that each entry E contains:\n",
        "      E['label']: (integer) the rumour veracity annotation.\n",
        "      E['source_text']: (string) the source tweet text.\n",
        "      E['reactions']: (List of strings) the texts from the tweet replies.\n",
        "  \"\"\"\n",
        "  labels_and_texts = []\n",
        "  for t in topics:\n",
        "    rumours = [\n",
        "      f.path for f in os.scandir(t + '/rumours') if f.is_dir()\n",
        "    ]\n",
        "    for r in rumours:\n",
        "      annotation, structure, source_tweets, reactions = load_pheme_data(r)\n",
        "      for source_tweet in source_tweets.values():\n",
        "        labels_and_texts.append({\n",
        "            'label' : convert_annotations(annotation, string = False),\n",
        "            'source_text' : source_tweet['text'],\n",
        "            'reactions' : [reaction['text'] for reaction in reactions.values()]\n",
        "        })\n",
        "\n",
        "  print('Read', len(labels_and_texts), 'annotated rumour tweet threads')\n",
        "\n",
        "  verified_labels_and_texts = list(\n",
        "      filter(lambda d : d['label'] != 2, labels_and_texts))\n",
        "\n",
        "  print('Returning', len(verified_labels_and_texts),\n",
        "        'verified rumour tweet threads')\n",
        "  \n",
        "  return verified_labels_and_texts"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8CMNokYICn0U",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "topics = [\n",
        "  f.path\n",
        "  for f in os.scandir('/tmp/PHEME/all-rnr-annotated-threads/')\n",
        "  if f.is_dir()\n",
        "]\n",
        "topics"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bLshZ9npwLuu",
        "colab_type": "text"
      },
      "source": [
        "## Load tokenizers, model architectures and pre-trained weights"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0Q-ehd0pbOsz",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "configuration = BertConfig()\n",
        "\n",
        "\"\"\"\n",
        "Uncomment any tokenizer-model pair for experimentation with other models.\n",
        "\"\"\"\n",
        "\n",
        "TOKENIZERS = {\n",
        "    # 'bert_base' : BertTokenizer.from_pretrained('bert-base-uncased'),\n",
        "    # 'bert_large' : BertTokenizer.from_pretrained('bert-large-uncased'),\n",
        "    # 'xlnet_base' : XLNetTokenizer.from_pretrained('xlnet-base-cased'),\n",
        "    # 'xlnet_large' : XLNetTokenizer.from_pretrained('xlnet-large-cased'),\n",
        "    # 'roberta_base' : RobertaTokenizer.from_pretrained('roberta-base'),\n",
        "    # 'roberta_large' : RobertaTokenizer.from_pretrained('roberta-large'),\n",
        "    'albert_base' : AlbertTokenizer.from_pretrained('albert-base-v2'),\n",
        "    # 'albert_large' : AlbertTokenizer.from_pretrained('albert-large-v2'),\n",
        "    # 'albert_xlarge' : AlbertTokenizer.from_pretrained('albert-xlarge-v2'),\n",
        "    # 'albert_xxlarge' : AlbertTokenizer.from_pretrained('albert-xxlarge-v2'),\n",
        "    # 't5_small' : T5Tokenizer.from_pretrained('t5-small'),\n",
        "    # 't5_base' : T5Tokenizer.from_pretrained('t5-base'),\n",
        "    # 't5_large' : T5Tokenizer.from_pretrained('t5-large'),\n",
        "    # 'electra_small' : ElectraTokenizer.from_pretrained('google/electra-small-discriminator'),\n",
        "    # 'electra_large' : ElectraTokenizer.from_pretrained('google/electra-large-discriminator'),\n",
        "}\n",
        "\n",
        "PRETRAINED_MODELS = {\n",
        "    # 'bert_base' : TFBertModel.from_pretrained('bert-base-uncased'),\n",
        "    # 'bert_large' : TFBertModel.from_pretrained('bert-large-uncased'),\n",
        "    # 'xlnet_base' : TFXLNetModel.from_pretrained('xlnet-base-cased'),\n",
        "    # 'xlnet_large' : TFXLNetModel.from_pretrained('xlnet-large-cased'),\n",
        "    # 'roberta_base' : TFRobertaModel.from_pretrained('roberta-base'),\n",
        "    # 'roberta_large' : TFRobertaModel.from_pretrained('roberta-large'),\n",
        "    'albert_base' : TFAlbertModel.from_pretrained('albert-base-v2'),\n",
        "    # 'albert_large' : TFAlbertModel.from_pretrained('albert-large-v2'),\n",
        "    # 'albert_xlarge' : TFAlbertModel.from_pretrained('albert-xlarge-v2'),\n",
        "    # 'albert_xxlarge' : TFAlbertModel.from_pretrained('albert-xxlarge-v2'),\n",
        "    # 't5_small' : TFT5Model.from_pretrained('t5-small'),\n",
        "    # 't5_base' : TFT5Model.from_pretrained('t5-base'),\n",
        "    # 't5_large' : TFT5Model.from_pretrained('t5-large'),\n",
        "    # 'electra_small' : TFElectraForPreTraining.from_pretrained('google/electra-small-discriminator'),\n",
        "    # 'electra_large' : TFElectraForPreTraining.from_pretrained('google/electra-large-discriminator')\n",
        "}\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KuI5IDx8xB2t",
        "colab_type": "text"
      },
      "source": [
        "## Tokenization and model inference functions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DROLZVZuDNbW",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Given that tweets are short, use a small value for reduced inference time\n",
        "TOKENIZER_MAX_SEQ_LENGTH = 64 \n",
        " \n",
        "\n",
        "def bert_embedding_model_inference(model, tokenizer, text):\n",
        "  \"\"\" Tokenizes and computes BERT `model` inference for input `text`. \"\"\"\n",
        "  input_ids = tf.constant(\n",
        "      tokenizer.encode(text, add_special_tokens=True))[None, :]  # Batch size 1\n",
        "  outputs = model(input_ids)\n",
        "  # The last hidden-state is the first element of the output tuple \n",
        "  last_hidden_states = outputs[0]\n",
        "  cls_embedding = last_hidden_states[0][0]\n",
        "  return cls_embedding\n",
        "\n",
        "\n",
        "def albert_embedding_model_inference(model, tokenizer, text):\n",
        "  \"\"\" Tokenizes and computes ALBERT `model` inference for input `text`. \"\"\"\n",
        "  encoded_text = tf.constant(tokenizer.encode(text))[None, :]\n",
        "  outputs = model(encoded_text)\n",
        "  # The last hidden-state is the first element of the output tuple \n",
        "  last_hidden_states = outputs[0]\n",
        "  cls_embedding = last_hidden_states[0][0]\n",
        "  return cls_embedding\n",
        "\n",
        "\n",
        "def albert_embedding_batch_model_inference(model, tokenizer, text_batch):\n",
        "  \"\"\" Tokenizes and computes ALBERT `model` inference for input `text_batch`.\n",
        "  \"\"\"\n",
        "  encoded_text = [\n",
        "      tf.constant(tokenizer.encode(\n",
        "          t, max_length=TOKENIZER_MAX_SEQ_LENGTH, pad_to_max_length=True))\n",
        "      for t in text_batch\n",
        "  ]\n",
        "  encoded_batch = tf.stack(encoded_text)\n",
        "  outputs = model(encoded_text)\n",
        "  cls_embeddings = []\n",
        "  for last_hidden_state in outputs[0]:\n",
        "    cls_embeddings.append(last_hidden_state[0])\n",
        "  return cls_embeddings\n",
        "\n",
        "\n",
        "def t5_embedding_model_inference(model, tokenizer, text):\n",
        "  \"\"\" Tokenizes and computes T5 `model` inference for input `text`. \"\"\"\n",
        "  inputs = tokenizer.encode(text, return_tensors=\"tf\")  # Batch size 1\n",
        "  outputs = model(inputs, decoder_input_ids=inputs)\n",
        "  # The last hidden-state is the first element of the output tuple \n",
        "  last_hidden_states = outputs[0]\n",
        "  cls_embedding = last_hidden_states[0][0]\n",
        "  return cls_embedding\n",
        "\n",
        "'''\n",
        "More model inference functions can be added, following documentation on\n",
        "https://huggingface.co/transformers/\n",
        "'''\n",
        "\n",
        "INFERENCE_MODELS = {\n",
        "  'bert' : bert_embedding_model_inference,\n",
        "  'albert' : albert_embedding_model_inference,\n",
        "  'albert_batch' : albert_embedding_batch_model_inference,\n",
        "  't5' : t5_embedding_model_inference\n",
        "}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kVX1oiSNxKnz",
        "colab_type": "text"
      },
      "source": [
        "## Testing tokenization and model inference with an example"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pXoyyToafH2T",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "PRETRAINED_MODEL = 'albert_base'\n",
        "INFERENCE = 'albert'\n",
        "\n",
        "test_inference = INFERENCE_MODELS[INFERENCE](\n",
        "    PRETRAINED_MODELS[PRETRAINED_MODEL],\n",
        "    TOKENIZERS[PRETRAINED_MODEL],\n",
        "    \"Why is the sky blue\"\n",
        ")\n",
        "\n",
        "print('Example embedding', test_inference)\n",
        "\n",
        "embedding_dim = test_inference.shape[0]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vzRd1HSexU-K",
        "colab_type": "text"
      },
      "source": [
        "## Hyperparameters for the classifier to be trained"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lTnxGZcWYsJP",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class HParams(object):\n",
        "  \"\"\"Hyperparameters used for training.\"\"\"\n",
        "  def __init__(self):\n",
        "    ### dataset parameters\n",
        "    self.num_classes = 2\n",
        "    self.embedding_dim = embedding_dim\n",
        "    ### neural graph learning parameters\n",
        "    self.distance_type = nsl.configs.DistanceType.L2\n",
        "    self.graph_regularization_multiplier = 0.1\n",
        "    self.num_neighbors = 5\n",
        "    ### model architecture\n",
        "    self.num_fc_units = [64, 64]\n",
        "    ### training parameters\n",
        "    self.train_epochs = 50\n",
        "    self.batch_size = 128\n",
        "    self.encoder_inference_batch_size = 32\n",
        "    self.dropout_rate = 0.2\n",
        "    ### eval parameters\n",
        "    self.eval_steps = None  # All instances in the test set are evaluated.\n",
        "\n",
        "HPARAMS = HParams()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zJ8cNaO0xhW5",
        "colab_type": "text"
      },
      "source": [
        "## Process dataset in batch or single inference mode"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cJtuP19ZRifL",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def add_batch_embeddings(labels_and_texts):\n",
        "  \"\"\"Splits `labels_and_texts` into batches and performs tokenization and\n",
        "  converts tweet text into its corresponding embedding.\n",
        "\n",
        "  Args:\n",
        "    labels_and_texts: A List of dictionaries such that each entry E contains:\n",
        "      E['label']: (integer) the rumour veracity annotation.\n",
        "      E['source_text']: (string) the source tweet text.\n",
        "      E['reactions']: (List of strings) the texts from the tweet replies.\n",
        "\n",
        "  Each entry E from labels_and_texts is updated, adding the following key,\n",
        "  value pairs:\n",
        "      E['source_embedding']: (Tensor of floats) embeddings for E['source_text']\n",
        "      E['reaction_embedding']: (List of float Tensors) embeddings for\n",
        "                               E['reactions'], up to HPARAMS.num_neighbors and\n",
        "                               in the corresponding E['reactions'] order.\n",
        "  \"\"\"\n",
        "  inputs = []\n",
        "  print('Accumulating inputs')\n",
        "  for e in labels_and_texts:\n",
        "    inputs.append(e['source_text'])\n",
        "    for r in e['reactions'][:HPARAMS.num_neighbors]:\n",
        "      # Alternative ways to select neighbors within a tweet thread can be used.\n",
        "      inputs.append(r)\n",
        "\n",
        "  def generate_batches(inputs, batch_size):\n",
        "    \"\"\"Splits `inputs` list into chunks of (up to) `batch_size` length.\"\"\"\n",
        "    for i in range(0, len(inputs), batch_size):\n",
        "      yield inputs[i: i + batch_size]\n",
        "\n",
        "  inferences = []\n",
        "  for i, batch in enumerate(generate_batches(\n",
        "      inputs, HPARAMS.encoder_inference_batch_size)):\n",
        "    print('Running model inference for batch', i)\n",
        "    batch_inference = INFERENCE_MODELS[INFERENCE](\n",
        "        PRETRAINED_MODELS[PRETRAINED_MODEL],\n",
        "        TOKENIZERS[PRETRAINED_MODEL],\n",
        "        batch)\n",
        "    for inference in batch_inference:\n",
        "      inferences.append(inference)\n",
        "\n",
        "  i = 0\n",
        "  for e in labels_and_texts:\n",
        "    e['source_embedding'] = inferences[i]\n",
        "    i += 1\n",
        "    e['reaction_embedding'] = []\n",
        "    for r in e['reactions'][:HPARAMS.num_neighbors]:\n",
        "      e['reaction_embedding'].append(inferences[i])\n",
        "      i += 1"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3PZsRpUqfIPx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def add_embeddings(labels_and_texts):\n",
        "  \"\"\"Performs tokenization and model inference for each element of\n",
        "  labels_and_texts, updating it with computed embeddings.\n",
        "\n",
        "  Args:\n",
        "    labels_and_texts: A List of dictionaries such that each entry E contains:\n",
        "      E['label']: (integer) the rumour veracity annotation.\n",
        "      E['source_text']: (string) the source tweet text.\n",
        "      E['reactions']: (List of strings) the texts from the tweet replies.\n",
        "\n",
        "  Each entry E from labels_and_texts is updated, adding the following key,\n",
        "  value pairs:\n",
        "      E['source_embedding']: (Tensor of floats) embeddings for E['source_text']\n",
        "      E['reaction_embedding']: (List of float Tensors) embeddings for\n",
        "                               E['reactions'], up to HPARAMS.num_neighbors and\n",
        "                               in the corresponding E['reactions'] order.\n",
        "  \"\"\"\n",
        "  for index, label_and_texts in enumerate(labels_and_texts):\n",
        "    if index % 100 == 0:\n",
        "      print('Computing embeddings for item', index)\n",
        "    label_and_texts['source_embedding'] = INFERENCE_MODELS[INFERENCE](\n",
        "        PRETRAINED_MODELS[PRETRAINED_MODEL],\n",
        "        TOKENIZERS[PRETRAINED_MODEL],\n",
        "        label_and_texts['source_text'])\n",
        "    label_and_texts['reaction_embedding'] = []\n",
        "    for r in label_and_texts['reactions'][:HPARAMS.num_neighbors]:\n",
        "      # Alternative ways to select neighbors within a tweet thread can be used.\n",
        "      label_and_texts['reaction_embedding'].append(INFERENCE_MODELS[INFERENCE](\n",
        "        PRETRAINED_MODELS[PRETRAINED_MODEL],\n",
        "        TOKENIZERS[PRETRAINED_MODEL],\n",
        "        r))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "toHnl4pqNsLz",
        "colab_type": "text"
      },
      "source": [
        "## Load and inspect dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "39ykV5MQ-mPX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "labels_and_texts = load_labels_and_texts(topics)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bmIafmeDOrLO",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "for e in random.sample(labels_and_texts, 3):\n",
        "  pprint.pprint(e)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L8TZRgKxPPiO",
        "colab_type": "text"
      },
      "source": [
        "## Compute and store textual embeddings"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wwO3F4knNvrb",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "if 'batch' in INFERENCE:\n",
        "  add_batch_embeddings(labels_and_texts)\n",
        "else:\n",
        "  add_embeddings(labels_and_texts)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5Wvfw7vSyWCb",
        "colab_type": "text"
      },
      "source": [
        "## Create train and test datasets\n",
        "\n",
        "Using a 80:20 train:test split"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tLv4SaQz7C-v",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Alternative ways to split include a split by time or by news event.\n",
        "random.shuffle(labels_and_texts)\n",
        "\n",
        "train_size = int(0.8 * len(labels_and_texts))\n",
        "\n",
        "TRAIN_DATA = labels_and_texts[:train_size]\n",
        "TEST_DATA = labels_and_texts[train_size:]\n",
        "\n",
        "# Constants used to identify neighbor features in the input.\n",
        "NBR_FEATURE_PREFIX = 'NL_nbr_'\n",
        "NBR_WEIGHT_SUFFIX = '_weight'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fOTeLCYE7ILJ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def create_np_tensors_from_datum(datum, propagate_label=True):\n",
        "  \"\"\"Creates a node and neighbor numpy tensors given input datum.\n",
        "\n",
        "  Args:\n",
        "    datum: A dictionary with node and neighbor features and annotation label.\n",
        "    propagate_label: Boolean indicating if we labels should be propagated to\n",
        "    neighbors.\n",
        "\n",
        "  Returns:\n",
        "    A tuple T (node_tensor, neighbor_tensors) such that:\n",
        "      T[0]: a dictionary containing node embeddings and label\n",
        "      T[1]: a List of dictionaries, each containing embeddings for a\n",
        "      neighbor and, if propagate_label is true, the corresponding label.\n",
        "\n",
        "  \"\"\"\n",
        "  label_tensor =  datum['label']\n",
        "  \n",
        "  def get_float32_tensor(d):\n",
        "    np_array = np.array(d, dtype='f')\n",
        "    return np_array\n",
        "    \n",
        "  node_tensor = {\n",
        "    'embedding' : get_float32_tensor(datum['source_embedding']),\n",
        "    'label' : label_tensor\n",
        "  }\n",
        "\n",
        "  neighbor_tensors = []\n",
        "  for re in datum['reaction_embedding']:\n",
        "    tensor = {\n",
        "      'embedding' : get_float32_tensor(re),\n",
        "    }\n",
        "    if propagate_label:\n",
        "      tensor['label'] = label_tensor\n",
        "    neighbor_tensors.append(dict(tensor))\n",
        "\n",
        "  return node_tensor, neighbor_tensors"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vHoegDArM5Tt",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "TRAIN_TENSORS = [\n",
        "  create_np_tensors_from_datum(d) for d in TRAIN_DATA\n",
        "]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tA2jPQT5NI61",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "TEST_TENSORS = [\n",
        "  create_np_tensors_from_datum(d) for d in TEST_DATA\n",
        "]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-rbriddOkZ76",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "TRAIN_TENSORS[0][0], TEST_TENSORS[0][0]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qA4_fjIn2Hcc",
        "colab_type": "text"
      },
      "source": [
        "## Create train and test tf.data.TFRecordDataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3V_ohmyk7P2E",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def make_dataset(tf_features, training=False):\n",
        "  \"\"\"Creates a `tf.data.TFRecordDataset`.\n",
        "\n",
        "  Args:\n",
        "    tf_features: List of (node_tensor, neighbor_tensors) tuples.\n",
        "    training: Boolean indicating if we are in training mode.\n",
        "\n",
        "  Returns:\n",
        "    An instance of `tf.data.TFRecordDataset` containing the `tf.train.Example`\n",
        "    objects.\n",
        "  \"\"\"\n",
        "\n",
        "  def get_tf_examples_with_nsl_signals(node_feature, neighbor_features):\n",
        "    \"\"\"Merges `neighbor_features` and `node_feature`.\n",
        "\n",
        "    Args:\n",
        "      node_feature: A dictionary of `tf.train.Feature`.\n",
        "      neighbor_features: A list of `tf.train.Feature` dictionaries.\n",
        "\n",
        "    Returns:\n",
        "      A pair whose first value is a dictionary containing relevant features\n",
        "      and whose second value contains the ground truth label.\n",
        "    \"\"\"\n",
        "    feature_dict = dict(node_feature)\n",
        "    # We also extract corresponding neighbor features in a similar manner to\n",
        "    # the features above during training.\n",
        "    if training:\n",
        "      for i in range(HPARAMS.num_neighbors):\n",
        "        nbr_feature_key = '{}{}_{}'.format(NBR_FEATURE_PREFIX, i, 'embedding')\n",
        "        nbr_weight_key = '{}{}{}'.format(NBR_FEATURE_PREFIX, i,\n",
        "                                         NBR_WEIGHT_SUFFIX)\n",
        "        if i < len(neighbor_features):\n",
        "          nf = neighbor_features[i]\n",
        "          feature_dict[nbr_feature_key] = nf['embedding']\n",
        "          feature_dict[nbr_weight_key] = 1.0\n",
        "        else:\n",
        "          feature_dict[nbr_feature_key] = np.zeros(\n",
        "              HPARAMS.embedding_dim, dtype='f')\n",
        "          feature_dict[nbr_weight_key] = 0.0\n",
        "\n",
        "    label = feature_dict.pop('label')\n",
        "    \n",
        "    return feature_dict, label\n",
        "\n",
        "  print('Adding NSL features for entries')\n",
        "  tensors_with_nsl = {}\n",
        "  labels = []\n",
        "  for (node, neighbors) in tf_features:\n",
        "    feature_dict, label = get_tf_examples_with_nsl_signals(node, neighbors)\n",
        "    for k,v in feature_dict.items():\n",
        "      if k not in tensors_with_nsl:\n",
        "        tensors_with_nsl[k] = []\n",
        "      tensors_with_nsl[k].append(v)\n",
        "    labels.append(label)\n",
        "  \n",
        "  print('Creating tf.data.Dataset from tensors')\n",
        "  dataset = tf.data.Dataset.from_tensor_slices((tensors_with_nsl, labels))\n",
        "  dataset = dataset.batch(HPARAMS.batch_size)\n",
        "  return dataset"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5aiVbfAHSk_G",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "train_dataset = make_dataset(TRAIN_TENSORS, training=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E0WuuQ4GPOsr",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "test_dataset = make_dataset(TEST_TENSORS)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mJQco49q2S8q",
        "colab_type": "text"
      },
      "source": [
        "## Inspecting train dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rWpE3J8E7TOS",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "for feature_batch, label_batch in train_dataset.take(1):\n",
        "  print('Feature list:', list(feature_batch.keys()))\n",
        "  print('Batch of inputs:', feature_batch['embedding'])\n",
        "  nbr_feature_key = '{}{}_{}'.format(NBR_FEATURE_PREFIX, 0, 'embedding')\n",
        "  nbr_weight_key = '{}{}{}'.format(NBR_FEATURE_PREFIX, 0, NBR_WEIGHT_SUFFIX)\n",
        "  print('Batch of neighbor inputs:', feature_batch[nbr_feature_key])\n",
        "  print('Batch of neighbor weights:',\n",
        "        tf.reshape(feature_batch[nbr_weight_key], [-1]))\n",
        "  print('Batch of labels:', label_batch)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7VXL7f4E2W9d",
        "colab_type": "text"
      },
      "source": [
        "## Inspecting test dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Y1RQUbEv7XAu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "for feature_batch, label_batch in test_dataset.take(1):\n",
        "  print('Feature list:', list(feature_batch.keys()))\n",
        "  print('Batch of inputs:', feature_batch['embedding'])\n",
        "  print('Batch of labels:', label_batch)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WE97Uc-f2cDK",
        "colab_type": "text"
      },
      "source": [
        "## Multi-layer perceptron classification model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dumju3JG2xHS",
        "colab_type": "text"
      },
      "source": [
        "### Sequential MLP model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "t-eTakWg7atT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def make_mlp_sequential_model(hparams):\n",
        "  \"\"\"Creates a sequential multi-layer perceptron model.\"\"\"\n",
        "  model = tf.keras.Sequential()\n",
        "  model.add(\n",
        "      tf.keras.layers.InputLayer(\n",
        "          input_shape=(hparams.embedding_dim,), name='embedding'))\n",
        "  for num_units in hparams.num_fc_units:\n",
        "    model.add(tf.keras.layers.Dense(num_units, activation='relu'))\n",
        "    # For sequential models, by default, Keras ensures that the 'dropout' layer\n",
        "    # is invoked only during training.\n",
        "    model.add(tf.keras.layers.Dropout(hparams.dropout_rate))\n",
        "  model.add(tf.keras.layers.Dense(hparams.num_classes, activation='softmax'))\n",
        "  return model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wZlS83MU21K8",
        "colab_type": "text"
      },
      "source": [
        "### Functional MLP model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sPf_awKm7eR2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def make_mlp_functional_model(hparams):\n",
        "  \"\"\"Creates a functional API-based multi-layer perceptron model.\"\"\"\n",
        "  inputs = tf.keras.Input(\n",
        "      shape=(hparams.embedding_dim,), dtype='int64', name='embedding')\n",
        "\n",
        "  cur_layer = inputs\n",
        "\n",
        "  for num_units in hparams.num_fc_units:\n",
        "    cur_layer = tf.keras.layers.Dense(num_units, activation='relu')(cur_layer)\n",
        "    # For functional models, by default, Keras ensures that the 'dropout' layer\n",
        "    # is invoked only during training.\n",
        "    cur_layer = tf.keras.layers.Dropout(hparams.dropout_rate)(cur_layer)\n",
        "\n",
        "  outputs = tf.keras.layers.Dense(\n",
        "      hparams.num_classes, activation='softmax')(\n",
        "          cur_layer)\n",
        "\n",
        "  model = tf.keras.Model(inputs, outputs=outputs)\n",
        "  return model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O2HxQOtS24_m",
        "colab_type": "text"
      },
      "source": [
        "### tf.Keras.Model subclass MLP"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xanEfOFy7hSP",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def make_mlp_subclass_model(hparams):\n",
        "  \"\"\"Creates a multi-layer perceptron subclass model in Keras.\"\"\"\n",
        "\n",
        "  class MLP(tf.keras.Model):\n",
        "    \"\"\"Subclass model defining a multi-layer perceptron.\"\"\"\n",
        "\n",
        "    def __init__(self):\n",
        "      super(MLP, self).__init__()\n",
        "      self.dense_layers = [\n",
        "          tf.keras.layers.Dense(num_units, activation='relu')\n",
        "          for num_units in hparams.num_fc_units\n",
        "      ]\n",
        "      self.dropout_layer = tf.keras.layers.Dropout(hparams.dropout_rate)\n",
        "      self.output_layer = tf.keras.layers.Dense(\n",
        "          hparams.num_classes, activation='softmax')\n",
        "\n",
        "    def call(self, inputs, training=False):\n",
        "      cur_layer = inputs['embedding']\n",
        "      for dense_layer in self.dense_layers:\n",
        "        cur_layer = dense_layer(cur_layer)\n",
        "        cur_layer = self.dropout_layer(cur_layer, training=training)\n",
        "\n",
        "      outputs = self.output_layer(cur_layer)\n",
        "\n",
        "      return outputs\n",
        "\n",
        "  return MLP()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z2goiRp42tGn",
        "colab_type": "text"
      },
      "source": [
        "# Base MLP model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5uW6jw8p7j7U",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Create a base MLP model using the functional API.\n",
        "# Alternatively, you can also create a sequential or subclass base model using\n",
        "# the make_mlp_sequential_model() or make_mlp_subclass_model() functions\n",
        "# respectively, defined above. Note that if a subclass model is used, its\n",
        "# summary cannot be generated until it is built.\n",
        "base_model_tag, base_model = 'FUNCTIONAL', make_mlp_functional_model(HPARAMS)\n",
        "base_model.summary()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q3pur3sN2nhx",
        "colab_type": "text"
      },
      "source": [
        "## Compile and train the base MLP model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "h9BexL637myp",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "base_model.compile(\n",
        "    optimizer='adam',\n",
        "    loss='sparse_categorical_crossentropy',\n",
        "    metrics=['accuracy'])\n",
        "base_model.fit(train_dataset, epochs=HPARAMS.train_epochs, verbose=1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J1QOMcbB2_yC",
        "colab_type": "text"
      },
      "source": [
        "## Evaluate base model on test dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cA544ann7tJF",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Helper function to print evaluation metrics.\n",
        "def print_metrics(model_desc, eval_metrics):\n",
        "  \"\"\"Prints evaluation metrics.\n",
        "\n",
        "  Args:\n",
        "    model_desc: A description of the model.\n",
        "    eval_metrics: A dictionary mapping metric names to corresponding values. It\n",
        "      must contain the loss and accuracy metrics.\n",
        "  \"\"\"\n",
        "  print('\\n')\n",
        "  print('Eval accuracy for ', model_desc, ': ', eval_metrics['accuracy'])\n",
        "  print('Eval loss for ', model_desc, ': ', eval_metrics['loss'])\n",
        "  if 'graph_loss' in eval_metrics:\n",
        "    print('Eval graph loss for ', model_desc, ': ', eval_metrics['graph_loss'])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ogYU8Qfe7ymg",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "eval_results = dict(\n",
        "    zip(base_model.metrics_names,\n",
        "        base_model.evaluate(test_dataset, steps=HPARAMS.eval_steps)))\n",
        "print_metrics('Base MLP model', eval_results)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G7K79yxz3DvH",
        "colab_type": "text"
      },
      "source": [
        "## Create, compile and train MLP model with graph regularization"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0ZW2FWPq74Ee",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Build a new base MLP model.\n",
        "base_reg_model_tag, base_reg_model = 'FUNCTIONAL', make_mlp_functional_model(\n",
        "    HPARAMS)\n",
        "\n",
        "# Wrap the base MLP model with graph regularization.\n",
        "graph_reg_config = nsl.configs.make_graph_reg_config(\n",
        "    max_neighbors=HPARAMS.num_neighbors,\n",
        "    multiplier=HPARAMS.graph_regularization_multiplier,\n",
        "    distance_type=HPARAMS.distance_type,\n",
        "    sum_over_axis=-1)\n",
        "graph_reg_model = nsl.keras.GraphRegularization(base_reg_model,\n",
        "                                                graph_reg_config)\n",
        "graph_reg_model.compile(\n",
        "    optimizer='adam',\n",
        "    loss='sparse_categorical_crossentropy',\n",
        "    metrics=['accuracy'])\n",
        "graph_reg_model.fit(train_dataset, epochs=HPARAMS.train_epochs, verbose=1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Zh2WuNom3J4C",
        "colab_type": "text"
      },
      "source": [
        "## Evaluate MLP model with graph regularization on test dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "G-woB5rm8Cix",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "eval_results = dict(\n",
        "    zip(graph_reg_model.metrics_names,\n",
        "        graph_reg_model.evaluate(test_dataset, steps=HPARAMS.eval_steps)))\n",
        "print_metrics('MLP + graph regularization', eval_results)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}